# res = [{
#     'input_path': WindowsPath('C:/Users/zhiyo/.paddlex/temp/tmpy2uydn99.png'),
#     'layout_result': {
#         'input_path': WindowsPath('C:/Users/zhiyo/.paddlex/temp/tmpy2uydn99.png'),
#         'boxes': [{
#             'cls_id': 0,
#             'label': 'image',
#             'score': 0.6926079988479614,
#             'coordinate': [351.09897, 12.47476, 558.2179, 278.6578]
#         }, {
#             'cls_id': 0,
#             'label': 'image',
#             'score': 0.5912043452262878,
#             'coordinate': [30.052397, 10.415057, 563.4621, 285.6357]
#         }]
#     },
#     'ocr_result': {
#         'input_path': WindowsPath('C:/Users/zhiyo/.paddlex/temp/tmpy2uydn99.png'),
#         'dt_polys': [array([
#             [57, 58],
#             [133, 56],
#             [134, 82],
#             [57, 84]
#         ], dtype=int16), array([
#             [121, 53],
#             [202, 53],
#             [202, 81],
#             [121, 81]
#         ], dtype=int16), array([
#             [57, 99],
#             [151, 99],
#             [151, 128],
#             [57, 128]
#         ], dtype=int16), array([
#             [181, 97],
#             [259, 100],
#             [258, 126],
#             [180, 123]
#         ], dtype=int16), array([
#             [59, 144],
#             [112, 144],
#             [112, 165],
#             [59, 165]
#         ], dtype=int16), array([
#             [105, 149],
#             [122, 149],
#             [122, 161],
#             [105, 161]
#         ], dtype=int16), array([
#             [122, 145],
#             [306, 143],
#             [307, 163],
#             [122, 166]
#         ], dtype=int16), array([
#             [60, 188],
#             [125, 188],
#             [125, 209],
#             [60, 209]
#         ], dtype=int16), array([
#             [119, 187],
#             [366, 187],
#             [366, 211],
#             [119, 211]
#         ], dtype=int16), array([
#             [120, 223],
#             [245, 223],
#             [245, 247],
#             [120, 247]
#         ], dtype=int16), array([
#             [63, 297],
#             [191, 297],
#             [191, 318],
#             [63, 318]
#         ], dtype=int16), array([
#             [212, 300],
#             [482, 300],
#             [482, 317],
#             [212, 317]
#         ], dtype=int16)],
#         'dt_scores': [0.9642157549064224, 0.861694492297995, 0.9490653996726358, 0.8522472701298413, 0.8631437408522568,
#                       0.7307912136778733, 0.8106032693328751, 0.8641561996047484, 0.8411793186921872, 0.786179191467901,
#                       0.7756745009922611, 0.8639687765093175],
#         'rec_text': ['姓名', '韦小宝', '性别男', '民族汉', '出生', '囍', '1654年12月20日', '住址',
#                      '北京市东城区景山前街4号', '紫禁城敬事房', '公民身份证号码', '11204416541220243X'],
#         'rec_score': [0.9988320469856262, 0.9897541403770447, 0.9991774559020996, 0.9972536563873291,
#                       0.9992415308952332, 0.025916164740920067, 0.9981616735458374, 0.9969428181648254,
#                       0.9962704181671143, 0.9972632527351379, 0.9935100674629211, 0.9961172938346863]
#     },
#     'table_ocr_result': [],
#     'table_result': [],
#     'layout_parsing_result': {
#         'input_path': WindowsPath('C:/Users/zhiyo/.paddlex/temp/tmpy2uydn99.png'),
#         'parsing_result': [{
#             'layout_bbox': array([121., 53., 202., 81.], dtype=float32),
#             'words in text block': '韦小宝',
#             'layout': 'double'
#         }, {
#             'layout_bbox': array([57., 56., 134., 84.], dtype=float32),
#             'words in text block': '姓名',
#             'layout': 'double'
#         }, {
#             'layout_bbox': array([180., 97., 259., 126.], dtype=float32),
#             'words in text block': '民族汉',
#             'layout': 'single'
#         }, {
#             'layout_bbox': array([57., 99., 151., 128.], dtype=float32),
#             'words in text block': '性别男',
#             'layout': 'double'
#         }, {
#             'layout_bbox': array([122., 143., 307., 166.], dtype=float32),
#             'words in text block': '1654年12月20日',
#             'layout': 'double'
#         }, {
#             'layout_bbox': array([59., 144., 112., 165.], dtype=float32),
#             'words in text block': '出生',
#             'layout': 'double'
#         }, {
#             'layout_bbox': array([105., 149., 122., 161.], dtype=float32),
#             'words in text block': '囍',
#             'layout': 'double'
#         }, {
#             'layout_bbox': array([119., 187., 366., 211.], dtype=float32),
#             'words in text block': '北京市东城区景山前街4号',
#             'layout': 'single'
#         }, {
#             'layout_bbox': array([60., 188., 125., 209.], dtype=float32),
#             'words in text block': '住址',
#             'layout': 'double'
#         }, {
#             'layout_bbox': array([120., 223., 245., 247.], dtype=float32),
#             'words in text block': '紫禁城敬事房',
#             'layout': 'double'
#         }, {
#             'layout_bbox': array([63., 297., 191., 318.], dtype=float32),
#             'words in text block': '公民身份证号码',
#             'layout': 'double'
#         }, {
#             'layout_bbox': array([212., 300., 482., 317.], dtype=float32),
#             'words in text block': '11204416541220243X',
#             'layout': 'single'
#         }]
#     },
#     'oricls_result': {},
#     'unwarp_result': {},
#     'curve_result': []
# }]
import base64
import io

from PIL import Image, ImageDraw
from paddlex import create_pipeline

pipeline = create_pipeline(
    pipeline="PP-ChatOCRv3-doc",
    llm_name="ernie-4.0",
    llm_params={"api_type": "aistudio", "access_token": "9b0c7f625a8badfbc19691472e336ec582ea04a4"},
    layout_model="./official_models/PicoDet-L_layout_3cls",  # 版面区域检测模块模型s
    table_model="./official_models/SLANet_plus",  # 表格识别模型
    # text_det_model="./official_models/PP-OCRv4_server_det",  # 文本检测模块模型 PP-OCRv4_server_det
    # text_rec_model="./official_models/PP-OCRv4_server_rec",  # 文本识别模块模型 PP-OCRv4_server_rec
    text_det_model="./official_models/PP-OCRv4_mobile_det",  # 文本检测模块模型 PP-OCRv4_server_det
    text_rec_model="./official_models/PP-OCRv4_mobile_rec",  # 文本识别模块模型 PP-OCRv4_server_rec
    seal_text_det_model="./official_models/PP-OCRv4_mobile_seal_det",  # 印章文本检测
    # doc_image_ori_cls_model="PP-LCNet_x1_0_doc_ori",  # 文档图像方向分类模块模型
    text_det_batch_size=1,  # 该参数只能是1
    text_rec_batch_size=5
)

visual_result, visual_info = pipeline.visual_predict("./image/27.png")


def base64_to_pil(base64_str):
    img_data = base64.b64decode(base64_str)
    img = Image.open(io.BytesIO(img_data))
    return img


def draw_text_box(img, dt_polys, rec_text, rec_score):
    texts = []
    polys = []
    scores = []
    draw = ImageDraw.Draw(img)
    for poly, txt, score in zip(dt_polys, rec_text, rec_score):
        if score > 0.9:
            top_left = tuple(poly[0])
            bottom_right = tuple(poly[2])
            draw.rectangle([top_left, bottom_right], outline='red', width=2)
            texts.append(txt)
            polys.append(poly)
            scores.append(score)
    # 显示图像
    img.show()
    # 将图像转换为 Base64 字符串
    buffered = io.BytesIO()
    img.save(buffered, format="PNG")  # 保存为 PNG 格式
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

    # 返回给前端的 JSON 响应
    response = {
        'image': img_str,
        'polys': polys,  # 将 polys 添加到响应中
        'texts': texts,  # 如果需要的话，也可以返回 texts
        'scores': scores  # 如果需要的话，也可以返回 scores
    }
    return response


for res in visual_result:
    res.save_to_img("./output")
    input_path = res['input_path']
    ocr_result = res['ocr_result']
    dt_polys = ocr_result['dt_polys']  # 边框坐标点
    rec_text = ocr_result['rec_text']  # 坐标点对应的文本
    rec_score = ocr_result['rec_score']  # 置信度
    layout_result = res['layout_result']
    layout_img = layout_result.img  # 原始图像
    ocr_img = ocr_result.img  # 2倍原始图像
    draw_text_box(layout_img, dt_polys, rec_text, rec_score)
